# -*- encoding: utf-8 -*-
'''
@File    :   xjtu_reader.py
@Time    :   2022/3/23 9:05
@Author  :   ZhangChaoYang
@Desc    :   把西安交通大学原始数据文件，经过各种转换（FFT、CWT等）变成numpy ndarray，存入文件。
'''

import os, sys

sys.path.insert(0, os.getcwd())
from data_process import mat_reader
import numpy as np
from absl import flags
from absl import app

FLAGS = flags.FLAGS

flags.DEFINE_string('source_data', "E:\zcy\数据集\XiJiao\Data\XJTU-SY_Bearing_Datasets", '原始数据存在目录')

WC = ["35Hz12kN", "37.5Hz11kN", "40Hz10kN"]  # 35Hz12kN表示转速35r/s，径向力12kN
fs = 25.6e3  # 采样频率25.6kHz


def concat_by_bearing():
    '''
    把同一个轴承的数据拼成一个文件
    :return:
    '''

    def concat(in_dir, out_file):
        path = os.path.dirname(out_file)
        if not os.path.exists(path):
            os.makedirs(path)
        with open(out_file, "w") as fout:
            for file in os.listdir(in_dir):
                with open(os.path.join(in_dir, file), encoding="utf-8") as fin:
                    fin.readline()
                    for line in fin:
                        arr = line.strip().split(",")
                        if len(arr) == 2:
                            fout.write(arr[0] + "\n")  # 只使用水平方向的振动信号

    root_dir = FLAGS.source_data
    for wc_dir in os.listdir(root_dir):
        for bearing_file in os.listdir(os.path.join(root_dir, wc_dir)):
            bearing_dir = os.path.join(root_dir, wc_dir, bearing_file)
            out_file = os.path.join("data", "XJTU", wc_dir, bearing_file + ".txt")
            concat(bearing_dir, out_file)


class XJTUReader(mat_reader.MatFileReader):
    def __init__(self, infile, trans1d="", trans2d="", work_load=""):
        super(XJTUReader, self).__init__(infile, "xjtu", trans1d, trans2d, work_load)

    def get_label(self):
        return 1

    def window_size(self):
        return 200

    def read1d(self):
        basename = self.get_basename()
        if not basename.endswith(".txt"):
            return None
        if self.work_load:
            if not os.path.dirname(self.infile).endswith(self.work_load):
                return None
        array = []
        with open(self.infile, encoding="utf-8") as fin:
            lines = fin.readlines()
            for line in lines:
                content = line.strip()
                if content:
                    array.append(float(content))
        x = self.trans1d(array, 1, fs)
        return x

    def read2d(self):
        basename = self.get_basename()
        if not basename.endswith(".txt"):
            return None
        if self.work_load:
            if not os.path.dirname(self.infile).endswith(self.work_load):
                return None
        array = []
        with open(self.infile, encoding="utf-8") as fin:
            lines = fin.readlines()
            for line in lines:
                content = line.strip()
                if content:
                    array.append(float(content))
        x = self.trans2d(array, 1, fs)
        return x


def main(argv):
    concat_by_bearing()
    for wc in WC:
        wc_dir = os.path.join("corpus", "XJTU", wc)
        for file in os.listdir(wc_dir):
            infile = os.path.join(wc_dir, file)
            reader = XJTUReader(infile, trans1d="stat", work_load=wc)
            array = reader.read1d()
            # 西交的数据只计算一些统计特征，将来只用于做趋势预测
            out_file = os.path.join("corpus", "1d", reader.name, "stat", wc, os.path.basename(infile)[:-4])
            path = os.path.dirname(out_file)
            if not os.path.exists(path):
                os.makedirs(path)
            np.save(out_file, array)


if __name__ == '__main__':
    app.run(main)

# python .\data_process\xjtu_reader.py --source_data E:\zcy\数据集\XiJiao\Data\XJTU-SY_Bearing_Datasets
